import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import os
import numpy as np
import forestplot as fp

# ====================== 配置 ======================
EXPERIMENT = "classification"
MODEL = "rf"
input_csv = f"results_subgroup/{EXPERIMENT}/{MODEL}/subgroup_point_estimates.csv"
output_dir = f"results_subgroup/{EXPERIMENT}/{MODEL}/visualizations"
os.makedirs(output_dir, exist_ok=True)

# ====================== 读取数据 ======================
df = pd.read_csv(input_csv)
metrics = ["acc", "f1", "precision", "recall", "roc_auc"]

# 创建数据框
plot_data_list = []

for metric in metrics:
    for _, row in df.iterrows():
        estimate = row[metric]
        se = row[f"{metric}_se"]
        lower = estimate - 1.96 * se
        upper = estimate + 1.96 * se
        
        plot_data_list.append({
            "metric": metric,
            "estimate": estimate,
            "lower": lower,
            "upper": upper,
            "subgroup": row["subgroup"]
        })

plot_data = pd.DataFrame(plot_data_list)

# ====================== 为每个指标创建单独的森林图 ======================
for metric in metrics:
    metric_data = plot_data[plot_data["metric"] == metric].copy()
    
    # 重置索引以确保正确显示
    metric_data = metric_data.reset_index(drop=True)
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # 使用正确的参数调用 forestplot
    fp.forestplot(
        metric_data,
        estimate='estimate',
        ll='lower',
        hl='upper',
        varlabel='subgroup',  # 必需的参数：指定亚组标签列
        ylabel='Subgroup',
        xlabel=f'{metric.upper()} Estimate (95% CI)',
        title=f'Forest Plot of {metric.upper()} by Subgroup',
        **{'color': 'blue', 'marker': 'D'}
    )
    
    plt.tight_layout()
    plt.savefig(f'{output_dir}/forestplot_{metric}.png', dpi=300, bbox_inches='tight')
    plt.savefig(f'{output_dir}/forestplot_{metric}.pdf', bbox_inches='tight')
    plt.close()

print('森林图已保存到:', output_dir)